import os
import psycopg2
import argparse
import sys
from configparser import ConfigParser
from psycopg2.extensions import ISOLATION_LEVEL_AUTOCOMMIT
import logging
from sqlalchemy import create_engine, MetaData, Table
from sqlalchemy.engine import reflection
from sqlalchemy import text
from datetime import datetime
import pathlib
import subprocess
import winreg

v = os.path.abspath(__file__)  # .py文件的绝对目录
log_dir = ''
if getattr(sys, 'frozen', False):
    log_dir = os.path.dirname(sys.executable)
else:
    log_dir = os.path.dirname(os.path.abspath(__file__))
# 设置日志级别为 INFO，并将日志记录到文件中
logging.basicConfig(filename=os.path.join(log_dir, 'db.log'),
                    level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')


def main():

    try:
        v = os.path.abspath(__file__)  # .py文件的绝对目录
        base_dir = ''
        if getattr(sys, 'frozen', False):
            exe_dir = os.path.dirname(sys.executable)
        else:
            exe_dir = os.path.dirname(os.path.abspath(__file__))
        logging.info('current_dir:'+exe_dir)
        print('current_dir:'+exe_dir)
        config = ConfigParser()

        # 读取INI文件
        config.read(os.path.join(exe_dir, 'db_config.ini'), encoding="utf-8")

        # 获取指定section的option值
        dbname = config.get('db', 'dbname')
        user = config.get('db', 'user')
        password = config.get('db', 'password')
        host = config.get('db', 'host')
        port = config.get('db', 'port')
        print(
            f'db_config.ini :\n' +
            '     [db]\n'
            f'    dbname={dbname}\n'
            f'    user={user}\n'
            f'    password={password}\n'
            f'    host={host}\n'
            f'    port={port}')

        # 解析命令行参数
        parser = argparse.ArgumentParser(description='Process some integers.')
        parser.add_argument('integers', metavar='N', type=int, nargs='*',
                            default=[],
                            help='an integer for the accumulator')
        args = parser.parse_args()
        logging.info(args.integers)
        print(args.integers)
        type = 0  # 0代表安装,1代表启动,2代表备份
        if len(args.integers) > 0:
            type = args.integers[0]
        if (type == 2):
            database_dump(exe_dir, dbname, user, password, host, port)
            return
        createdb(dbname, user, password, host, port)
        items = os.listdir(os.path.join(exe_dir, "sql"))
        sorted_items = sorted(items, key=lambda x: x.lower())
        # 查找并执行.sql文件
        for file in sorted_items:
            if file.endswith(".sql"):
                excetsql(type, exe_dir, file, dbname,
                         user, password, host, port)

    except Exception as e:
        logging.error(e)
        print(e)


def createdb(dbname, user, password, host, port):
    try:
        conn0 = psycopg2.connect(
            user=user,
            password=password,
            host=host,
            port=port,
        )
        conn0.set_isolation_level(ISOLATION_LEVEL_AUTOCOMMIT)
        # 创建cursor对象
        cur0 = conn0.cursor()
        create_db_sql = f"CREATE DATABASE {dbname} owner {user} encoding 'utf-8';"

        logging.info(create_db_sql)
        # 执行SQL命令创建数据库
        cur0.execute(create_db_sql)
        logging.info(
            f"CREATE DATABASE {dbname} owner {user} encoding 'utf-8';")
        print(f"successfully created database {dbname}.")
        logging.info(f"successfully created database {dbname}.")
        print(f"successfully created database {dbname}.")
        # 关闭cursor和连接
        cur0.close()
        conn0.close()
    except Exception as e:
        logging.error(e)
        print(e)
    finally:
        # 释放数据库连接
        try:
            if conn0 is not None:
                conn0.close()
                logging.info("connect is closed")
        except psycopg2.Error as e:
            logging.info("connect  close error:", e)
            print("connect  close error:", e)


def excetsql(type, exe_dir, file, dbname, user, password, host, port):
    # PostgreSQL数据库连接信息
    try:
        conn = psycopg2.connect(
            dbname=dbname,
            user=user,
            password=password,
            host=host,
            port=port,
        )
        initfile = file.startswith("00")
        logging.info("find sql file:"+file)
        print("find sql file: "+file)

        if (type == 0 or (type == 1 and initfile == False)):
            logging.info("excete sql :"+file)
            print("excete sql :"+file)
            with open(os.path.join(exe_dir, 'sql', file), 'r', encoding="utf-8") as f:
                sql_script = f.read()
                # 执行SQL脚本
                with conn.cursor() as cursor:
                    cursor.execute(sql_script)
                    conn.commit()
                    print(f"{file} executed.")
                    logging.info(f"{file} executed.")
            f.close()
    except Exception as e:
        logging.error(e)
        print(e)
    finally:
        # 释放数据库连接
        try:
            if conn is not None:
                conn.close()
                logging.info("connect is closed")
                print("connect is closed")
        except psycopg2.Error as e:
            logging.info("connect  close error:", e)
            print("connect  close error:", e)


def database_dump(exe_dir, dbname, user, password, host, port):
    try:

        # 获取当前日期时间
        now = datetime.now()
        # 使用strftime进行格式化
        formatted_date = now.strftime('%Y-%m-%d-%H-%M-%S')
        back_dir = os.path.join(exe_dir, "db_back")

        print(f"DDL dir:{back_dir}")

        back_file = os.path.join(back_dir, f"{dbname}-{formatted_date}.sql")
        result = backup_pg_dump(host, port, user, password, dbname,
                                os.path.join(exe_dir, back_file))
        if (result is True):
            return
        engine = create_engine(
            f'postgresql://{user}:{password}@{host}:{port}/{dbname}')
        # 获取inspector对象
        # inspector = reflection.Inspector.from_engine(engine)
        # 查询数据库中所有表的名称
        # 连接到 Postgres 数据库
        conn = engine.connect()
        # 获取元数据
        metadata = MetaData()
        metadata.reflect(bind=engine)

        # 创建文件夹如果不存在，则新建一个文件夹
        pathlib.Path(back_dir).mkdir(parents=True, exist_ok=True)

        with open(back_file, 'a', encoding="utf-8") as f:
            f.write("SET statement_timeout = 0;" + "\n")
            f.write("SET lock_timeout = 0;" + "\n")
            f.write("SET idle_in_transaction_session_timeout = 0;" + "\n")
            f.write("SET client_encoding = 'UTF8';" + "\n")
            f.write("SET standard_conforming_strings = on;" + "\n")
            f.write("SELECT pg_catalog.set_config('search_path', '', false);" + "\n")
            f.write("SET check_function_bodies = false;" + "\n")
            f.write("SET xmloption = content;")
            f.write("SET client_min_messages = warning;" + "\n")
            f.write("SET row_security = off;" + "\n")
            f.write("SET default_tablespace = '';" + "\n")
            f.write("SET default_table_access_method = heap;" + "\n")
            f.write("\n\n")

            # 遍历数据库中的每个表
            for table in metadata.sorted_tables:
                print(f"{table.schema}.{table.name}")
                table_schema = table.schema
                if (table.schema is None):
                    table_schema = "public"

                f.write("-- "+str(table))
                # 获取表的注释
                comment = conn.execute(
                    text(f"SELECT obj_description('{table_schema}.{table.name}'::regclass)")).fetchone()[0]
                if (comment is not None):
                    f.write(
                        f"  {comment} ;\n")
                f.write('\n')
                # f.write(metadata.get_create_table_stmt(table))
                # 获取表的建表语句
                create_statement = []
                for column in table.columns:
                    column_type = ''
                    if (column.type._is_array):
                        type = column.type
                        item_type = getattr(type, 'item_type', None)
                        column_type += f"{item_type}[]"
                    else:
                        column_type += str(column.type)

                    if column.primary_key:
                        column_type += " PRIMARY KEY"
                    if column.default is not None:
                        column_type += " DEFAULT '{0}'".format(column.default)
                    if column.nullable is False:
                        column_type += " NOT NULL"
                    create_statement.append(f"{column.name} {column_type}")

                f.write(
                    f"CREATE TABLE if not exists  {table_schema}.{table.name} ({', '.join(create_statement)}) ;\n\n")
                f.write(
                    f"ALTER TABLE {table_schema}.{table.name} OWNER TO {user} ;\n")

                if (comment is not None):
                    f.write(
                        f"comment on table {table_schema}.{table.name} is '{comment}' ;\n")
                # 获取表的索引信息
                indexes = conn.execute(text(
                    f"SELECT indexname, indexdef FROM pg_indexes WHERE tablename = '{table.name}'"))
                for index in indexes:
                    # print(f"索引名称: {index[0]}")
                    findresult = index[1].find('UNIQUE')
                    if findresult == -1:
                        f.write(f"{index[1]} ;\n")

                col_index = 0
                for column in table.columns:
                    # print(f"字段名: {column.name}")
                    # 获取字段注释
                    comment = conn.execute(text(
                        f"SELECT col_description('{table_schema}.{table.name}'::regclass, {col_index})")).fetchone()[0]
                    if (comment is not None):
                        f.write(
                            f"comment on column {table_schema}.{table.name}.{column.name} is '{comment}' ;\n")
                    col_index += 1
                f.write("\n\n\n\n")
                # 关闭连接
            conn.close()

    except Exception as e:
        logging.error(e)
        print(e)
    finally:
        # 释放数据库连接
        try:
            if f is not None:
                f.close()
            if conn is not None:
                conn.close()
                logging.info("connect is closed")
                print("connect is closed")
        except psycopg2.Error as e:
            logging.info("connect  close error:", e)
            print("connect  close error:", e)


def backup_pg_dump(host, port, username, password, database_name, backup_file):
    try:
        key = winreg.OpenKey(winreg.HKEY_LOCAL_MACHINE,
                             r'SOFTWARE\Microsoft\Windows\CurrentVersion\Uninstall\PostgreSQL 12')
        value, _ = winreg.QueryValueEx(key, 'InstallLocation')
        winreg.CloseKey(key)
        if value == '':
            print('pg_dump not found')
            logging.info('pg_dump not found')
            return False
        print(value)
        cmds = [os.path.join(value, r'bin\pg_dump'), '-h', host, '-U',
                username, '-p', port, '-W',  '-f', backup_file, '-d', database_name]
        cmds.append('-s')
        # 使用 subprocess 模块执行 pg_dump 命令备份数据库
        subprocess.run(cmds)
        print('pg_dump success')
        return True
    except Exception as e:
        logging.error('pg_dump fail, use custom backup:' + e)
        print(e)
        return False


if __name__ == "__main__":
    main()
